import sys
sys.path.append('/home/llm_user/index/meta-learning/stable_meta_learning')
from stable_baselines3.common.vec_env import DummyVecEnv

import torch

from stable_meta_learning.pearl import PEARL_SAC
from stable_meta_learning.envs import make_env
from stable_meta_learning.pearl import MultiTaskHerReplayBuffer

train_env_1 = make_env('PandaPickAndPlace-v3',mass=0.5)
train_env_2 = make_env('PandaPickAndPlace-v3',mass=1)
train_env_3 = make_env('PandaPickAndPlace-v3',mass=1.5)
train_env_4 = make_env('PandaPickAndPlace-v3',mass=2)
# test_env_1 = make_env('PandaPickAndPlace-v3',mass=10.0)

train_env = DummyVecEnv([train_env_1,train_env_2,train_env_3,train_env_4])
# test_env = DummyVecEnv([test_env_1])
model = PEARL_SAC.load('/home/llm_user/index/meta-learning/stable_meta_learning/checkpoints/PEARL_SAC3.zip',
                        env=train_env,
                        device=torch.device('cuda:0'),
)
model.learn(total_timesteps=2_000_000,progress_bar=True)
model.save('checkpoints/PEARL_SAC_CONTINUE')

train_env.close()